import httpx
from typing import List, Optional
from app.core.config import settings
import logging

logger = logging.getLogger(__name__)

class OllamaService:
    """Ollama服务管理"""
    
    def __init__(self, base_url: Optional[str] = None):
        self.base_url = base_url or settings.OLLAMA_BASE_URL
        self.timeout = 10.0
    
    async def test_connection(self) -> dict:
        """
        测试Ollama连接
        
        Returns:
            dict: 连接测试结果
        """
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                # 测试基本连接
                response = await client.get(f"{self.base_url}/api/tags")
                
                if response.status_code == 200:
                    data = response.json()
                    models = [model["name"] for model in data.get("models", [])]
                    
                    return {
                        "success": True,
                        "message": "Ollama连接成功",
                        "ollama_version": "OK",
                        "available_models": models
                    }
                else:
                    return {
                        "success": False,
                        "message": f"连接失败: HTTP {response.status_code}",
                        "ollama_version": None,
                        "available_models": None
                    }
        except httpx.ConnectError:
            logger.error(f"无法连接到Ollama服务: {self.base_url}")
            return {
                "success": False,
                "message": f"无法连接到Ollama服务 ({self.base_url})",
                "ollama_version": None,
                "available_models": None
            }
        except Exception as e:
            logger.error(f"Ollama连接测试失败: {str(e)}")
            return {
                "success": False,
                "message": f"连接测试失败: {str(e)}",
                "ollama_version": None,
                "available_models": None
            }
    
    async def list_models(self) -> List[dict]:
        """
        获取Ollama模型列表
        
        Returns:
            List[dict]: 模型列表
        """
        try:
            async with httpx.AsyncClient(timeout=self.timeout) as client:
                response = await client.get(f"{self.base_url}/api/tags")
                
                if response.status_code == 200:
                    data = response.json()
                    return data.get("models", [])
                else:
                    logger.error(f"获取模型列表失败: HTTP {response.status_code}")
                    return []
        except Exception as e:
            logger.error(f"获取Ollama模型列表失败: {str(e)}")
            return []
    
    async def pull_model(self, model_name: str):
        """
        拉取Ollama模型
        
        Args:
            model_name: 模型名称
        """
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                response = await client.post(
                    f"{self.base_url}/api/pull",
                    json={"name": model_name}
                )
                return response.status_code == 200
        except Exception as e:
            logger.error(f"拉取模型失败: {str(e)}")
            return False

# 全局实例
ollama_service = OllamaService()

